#! /usr/bin/python
"""Music sorter written in Python

Sorts music files into appropriate folders and renames them according to tagdata.

Usage: sortmusik [options] [root folder]

Options:
    -c, --cleanup   remove empty directories after processing music files
    -h, --help      show this help
    -k, --keep      don't touch music files with no artist or title data
    -v, --verbose   show verbose output
    --version       print version information

Examples:
    sortmusik               recursively sort all music in the current working folder into folders based on artist and album name
    sortmusik /var/music    sort music in /var/music
    sortmusik -c            sort music in /var/music and remove empty folders when done

Todo:
    read actual tagsize instead of just using 2k bytes
    add more tag-types
    add guesswork to guess tag-type (based on file extension)
    extend to include more versions of id3
    add an option to sort into directories based on first letter of artist name
    add an option to sort into directories based on genre
    split tag parsing in modules
    """

from time import time
import os
from os.path import join, isdir, islink, isfile
from shutil import move
import struct
import getopt
import sys
from UserDict import UserDict

global version
version = 0.1a

class WrongTagScheme(Exception):
    def __init__(self, value):
        self.value = value

    def __str__(self):
        return repr(self.value)

class TagFile:
    def __init__(self):
        self.tags = {"Artist": "", "Album": "", "Title": ""}

class ID3File(TagFile):
    """Class describing a file containing ID3 tags"""
    tag_identifiers = {"Artist": ["TPE1", "TPE2"],
                       "Album": ["TALB"],
                       "Title": ["TIT2"]}
    
    def __init__(self, data):
        TagFile.__init__(self)
        header = struct.unpack("3sbxxbbbb", data[:10])
        if header[0] == 'ID3' and \
           header[1] < 0xff and \
           header[2] < 0x80 and \
           header[3] < 0x80 and \
           header[4] < 0x80 and \
           header[5] < 0x80:
            self.get_id3tags(data)
        else:
            raise WrongTagScheme("ID3")

    def get_id3tags(self, data):
        """Finds the tags in the file data"""
        for key, values in ID3File.tag_identifiers.iteritems():
            for value in values:
                offset = data.find(value)
                if offset != -1:
                    offset = offset+len(value) + 3
                    (size,) = struct.unpack("b", data[offset])
                    offset += 4
                    self.tags[key] = data[offset:offset+size-1]
                    break

def yes_no(question):
    response = raw_input(question)
    lresp = response.lower()
    if lresp in ("y", "yes", "j", "ja"):
        return True
    else:
        return False

def strip_nulls(string):
	list = []
	for x in string:
		if x != "\x00" and x != "\xff" and x != "\xfe":
			list.append(x)
	return "".join(list)


def get_name(tags, root, filename, instance=-1):
    """Generate a filename for a file based on tagdata"""
    if instance == -1:
        path = strip_nulls(join(root, tags["Artist"], tags["Album"], tags["Title"]) + ".mp3")
    else:
        path = strip_nulls(join(root, tags["Artist"], tags["Album"], tags["Title"]) + "(" + str(instance) + ").mp3")

    try:
        if isfile(path) and path != filename:
            instance += 1
            return get_name(tags, root, filename, instance)
        else:
            return path
    except TypeError, exception_instance:
        print
        print "---- ERROR ----"
        print "Path was: " + path
        print exception_instance
        print sys.exc_info()[0]
        raise
            
def arrange_file(filename, tags, root):
    """Moves files around according to tagdata."""
    if keep:
        if not tags["Artist"] and not tags["Title"]:
            print "Keeping " + filename
            return
        
    if not tags["Artist"]:
        tags["Artist"] = "Unknown Artists"
    #if not tags["Album"]:
        #tags["Album"] = "Unknown Album"
    if not tags["Title"]:
        tags["Title"] = "Untitled Track"
    path = get_name(tags, root, filename)
    if path != filename:
        try:
            os.makedirs(os.path.dirname(path))
        except OSError:
            pass

        if not os.path.exists(path):
            move(filename, path)
        else:
            print "Something has gone wrong, aborting."
            sys.exit(2)
    else:
        if verbose:
            print "File already in place."

def parse_file(filename, root):
    """This function reads the file-header and passes the data to the appropriate class"""
    if os.stat(filename).st_size >= 10:
        if verbose:
            print "Parsing " + filename
        try:
            fsock = open(join(root, filename), "rb")
            data = fsock.read(2000)
            fsock.close()
            try:
                arrange_file(filename, ID3File(data).tags, root)
            except WrongTagScheme, scheme:
                if verbose:
                    print "The file was not tagged using " + scheme.value + "."
        except IOError:
            if verbose:
                print "File did not exist: " + filename
    else:
        if verbose:
            print "Skipping " + filename + ". File too small."

def traverse(dir, root=""):
    """This is the recursive funktion which traverses the root and passes filenames too parse_file"""
    if verbose:
        print "Traversing " + dir
    total_files = 0
    if not root:
        root = dir

    nodes = os.listdir(dir)
    for node in nodes:
        node = join(dir, node)
        if isdir(node) and not islink(node):
            total_files += traverse(node, root)
        elif isfile(node):
            try:
                parse_file(node, root)
                total_files += 1
            except SyntaxError:
                pass
        else:
            pass

    return total_files

def cleanup(dir):
    """Removes empty directories in the root. Called if the 'clean' flag is set"""
    dirs = 0
    nodes = os.listdir(dir)
    if nodes == []:
        os.rmdir(dir)
        dirs += 1
        if verbose:
            print "The directory " + dir + " was empty, and has been removed."
    else:
        for node in nodes:
            node = join(dir, node)
            if isdir(node) and not islink(node):
                foo = cleanup(node)     
                if foo:
                    dirs += cleanup(dir) + foo
    return dirs

def print_usage():
    """Prints usage"""
    print __doc__

def set_options(opts):
    global clean
    clean = 0
    global verbose
    verbose = 0
    global keep
    keep = 1
    global just_do_it
    just_do_it = 0
    """Parses options passed to the script and sets globals accordingly."""
    for o, a in opts:
            if o in ("-c", "--cleanup"):
                clean = 1
            elif o in ("-v", "--verbose"):
                print "Verbose"
                verbose = 1
            elif o in ("-k", "--nokeep"):
                keep = 0
            elif o in ("-y"):
                just_do_it = 1
            elif o in ("-h", "--help", "--usage"):
                print_usage()
                sys.exit()
            elif o in ("--version"):
                global version
                print "sortmusik " + str(version)
                sys.exit()

def main(argv):
    """Main function"""
    t1 = time()
    root = os.getcwd()
    try:
        opts, args = getopt.getopt(argv, "cvkhy", ["version", "help", "cleanup", "verbose", "nokeep"])
        try:
            root = args[0]
        except IndexError:
            pass

        set_options(opts)

    except getopt.GetoptError:
        pass

    if not just_do_it:
        if not yes_no("This may move files as well as create and delete directories\nAre you sure you want sortmusik to sort the files located in " + root + "? (y/N): "):
            sys.exit(2)

    response = traverse(root)

    if clean:
        if verbose:
            print "Performing cleanup"
        cleaned_dirs = cleanup(root)

    t2 = time()
    if verbose:
        print
    print str(response) + " files parsed in " + str(t2 - t1) + " seconds."
    print "(" + str(response / (t2 - t1)) + " files pr. second)"
    if clean:
        print str(cleaned_dirs) + " empty directories removed."

if __name__ == "__main__":
    main(sys.argv[1:])
    
